from PIL import Image, ImageEnhance, ImageOps
import pytesseract
import time

# 二值化算法
def binarizing(img,threshold):
    pixdata = img.load()
    w, h = img.size
    for y in range(h):
        for x in range(w):
            if pixdata[x, y] < threshold:
                pixdata[x, y] = 0
            else:
                pixdata[x, y] = 255
    return img


image = Image.open('./screenshot.png')

t1 = time.time()
# 图片一分为二
question = image.crop((18, 591, 1058, 865))
choices = image.crop((250, 969, 829, 1684))

t2 = time.time()
# 图片增强
question = ImageEnhance.Contrast(question).enhance(2)
choices = ImageEnhance.Contrast(choices).enhance(2)

t3 = time.time()
# 转化为灰度图
question = question.convert('L')
choices = choices.convert('L')

t4 = time.time()
# 把图片变成二值图像
question = binarizing(question, 190)
choices = binarizing(choices, 190)

t5 = time.time()
# 看一下处理后的图片是什么样
question.save('question3.png')
choices.save('choices3.png')

t6 = time.time()
# question图片转化文字
tessdata_dir_config = '--tessdata-dir "C:\\Program Files (x86)\\Tesseract-OCR\\tessdata" --psm 6'
text = pytesseract.image_to_string(question, lang='chi_sim', config=tessdata_dir_config)
text = text.replace('\n', '')
text = text.replace(' ', '')
print(text)
t7 = time.time()

# choices图片转化文字
text2 = pytesseract.image_to_string(choices, lang='chi_sim', config=tessdata_dir_config)
text_list = text2.split('\n一\n')
print(text_list)
text_list.insert(0, text)
t8 = time.time()

print('1', t2-t1)
print('2', t3-t2)
print('3', t4-t3)
print('4', t5-t4)
print('5', t6-t5)
print('6', t7-t6)
print('7', t8-t7)